from langchain.agents import initialize_agent, Tool
from langchain.agents import AgentType
from langchain_community.chat_models import ChatOpenAI
from langchain.tools import BaseTool
import os
from dotenv import load_dotenv
load_dotenv()  # 加载.env文件中的环境变量

import sqlite3
from typing import Optional, Type
from contextlib import contextmanager


# SQLite数据库连接上下文管理器
@contextmanager
def sqlite_connection(db_path: str):
    conn = None
    try:
        conn = sqlite3.connect(db_path)
        yield conn
        conn.commit()
    except Exception as e:
        if conn: conn.rollback()
        raise Exception(f"Database error: {str(e)}")
    finally:
        if conn: conn.close()

# 本地目录管理工具类
class DirectoryManagementTool(BaseTool):
    name = "directory_manager"
    description = "用于管理本地文件系统目录的工具，可以列出目录内容、创建目录和文件、读取文件内容"

    def _run(self, action: str, path: str, content: Optional[str] = None) -> str:
        if action == "list":
            if not os.path.exists(path):
                return f"路径不存在: {path}"
            if os.path.isfile(path):
                return f"{path} 是文件，不是目录"
            return str(os.listdir(path))
        elif action == "create_dir":
            os.makedirs(path, exist_ok=True)
            return f"目录创建成功: {path}"
        elif action == "create_file":
            if not content: raise ValueError("创建文件需要提供content参数")
            with open(path, 'w', encoding='utf-8') as f:
                f.write(content)
            return f"文件创建成功: {path}"
        elif action == "read_file":
            if not os.path.isfile(path):
                return f"{path} 不是文件或不存在"
            with open(path, 'r', encoding='utf-8') as f:
                return f.read()
        else:
            return f"不支持的操作: {action}"

    async def _arun(self, action: str, path: str, content: Optional[str] = None) -> str:
        raise NotImplementedError("此工具不支持异步调用")

# SQLite数据库工具类
class SQLiteTool(BaseTool):
    name = "sqlite_manager"
    description = "用于操作SQLite数据库的工具，需要传入SQL查询语句作为第一个参数，可选数据库路径作为第二个参数（默认为'notes_database.db'）。示例：query='SELECT * FROM notes', db_path='notes_database.db'"

    def _run(self, query: str, db_path: str = './notes_database.db') -> str:
        with sqlite_connection(db_path) as conn:
            cursor = conn.cursor()
            cursor.execute(query)
            if query.strip().lower().startswith('select'):
                columns = [desc[0] for desc in cursor.description]
                results = cursor.fetchall()
                return str([dict(zip(columns, row)) for row in results])
            return f"查询执行成功，影响行数: {cursor.rowcount}"

    async def _arun(self, db_path: str, query: str) -> str:
        raise NotImplementedError("此工具不支持异步调用")

# 初始化目录管理MCP代理
def init_directory_agent():
    # 注意：实际使用时需要设置DEEPSEEK_API_KEY环境变量
    llm = ChatOpenAI(
        temperature=0,
        model="deepseek-chat",
        base_url="https://api.deepseek.com/v1",
        api_key=os.getenv("DEEPSEEK_API_KEY")
    )
    tools = [DirectoryManagementTool()]
    return initialize_agent(
        tools, llm, agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True
    )

# 初始化数据库操作MCP代理
def init_database_agent():
    # 注意：实际使用时需要设置DEEPSEEK_API_KEY环境变量
    llm = ChatOpenAI(
        temperature=0,
        model="deepseek-chat",
        base_url="https://api.deepseek.com/v1",
        api_key=os.getenv("DEEPSEEK_API_KEY")
    )
    tools = [SQLiteTool()]
    return initialize_agent(
        tools, llm, agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True
    )

# 实例化两个MCP代理
directory_agent = init_directory_agent()
database_agent = init_database_agent()